#!/usr/bin/env python
# -*- coding: UTF-8 -*-
"""
@Project ：python_learning 
@File ：DatabaseManager.py
@IDE  ：PyCharm 
@Author ：李涵彬
@Date ：2024/5/29 下午2:25 
"""

import os
from dotenv import load_dotenv
import pandas as pd
from typing import Dict, Any, Optional
import logging
import contextlib
from sqlalchemy import create_engine, MetaData, text, Table
from sqlalchemy.orm import sessionmaker
from sqlalchemy.exc import InvalidRequestError, SQLAlchemyError

# 加载环境变量
load_dotenv()

# 日志配置
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


class DatabaseManager:
	SUPPORTED_DB_TYPES = {'mysql', 'postgresql', 'oracle', 'mssql'}

	def __init__(self, db_type: str, username: str = None, password: str = None, host: str = None,
				 port: int = None, database: str = None):
		"""
		初始化数据库连接器。

		参数:
		db_type (str): 数据库类型。
		username (str, 可选): 连接数据库所需的用户名，默认为 None。
		password (str, 可选): 连接数据库所需的密码，默认为 None。
		host (str, 可选): 数据库主机地址，默认为 None。
		port (int, 可选): 数据库监听端口，默认为 None。
		database (str, 可选): 要连接的数据库名称，默认为 None。

		无返回值。
		"""
		# 验证数据库类型的有效性
		self._validate_db_type(db_type)
		# 格式化数据库类型为小写
		self.db_type = db_type.lower()
		# 获取配置好的用户名，优先使用传入的用户名
		self.username = self._get_configured_value('username', username)
		# 获取配置好的密码，优先使用传入的密码
		self.password = self._get_configured_value('password', password)
		# 获取配置好的主机地址，优先使用传入的主机地址
		self.host = self._get_configured_value('host', host)
		# 验证并获取端口号
		self.port = self._validate_and_get_port(port)
		# 获取配置好的数据库名称，优先使用传入的数据库名称
		self.database = self._get_configured_value('database', database)

		# 初始化数据库连接
		self._init_db_connection()

	def _validate_db_type(self, db_type):
		"""
		验证给定的数据库类型是否为支持的类型。

		参数:
		- db_type: 要验证的数据库类型，为字符串。

		返回值:
		- 无返回值，但是如果给定的数据库类型不在支持的类型列表中，则抛出 ValueError。
		"""
		if db_type not in self.SUPPORTED_DB_TYPES:
			# 如果 db_type 不在支持的数据库类型列表中，则抛出 ValueError
			raise ValueError(
				f"Unsupported database type: {db_type}. Supported types are: {', '.join(self.SUPPORTED_DB_TYPES)}")

	def _get_configured_value(self, key, default_value):
		"""
		获取配置的值。

		从环境变量中尝试获取以'DB_'开头、紧接着是key（大写）的环境变量值。如果未找到且default_value为None，则返回环境变量的默认值。
		如果default_value不为None，则直接返回default_value。

		参数:
		key: str - 配置项的键名。
		default_value: 任意类型 - 当未找到配置项时的默认返回值。如果为None，则尝试从环境变量获取默认值。

		返回:
		任意类型 - 配置项的值或默认值。
		"""
		# 尝试从环境变量获取配置值，如果未设置且default_value为None，则返回默认值
		return os.getenv(f'DB_{key.upper()}', default_value) if default_value is None else default_value

	def _validate_and_get_port(self, port):
		"""
		验证并获取数据库端口号。

		首先检查传入的port参数是否为None，如果是，则尝试从环境变量'DB_PORT'中获取端口号，
		若环境变量未设置，则默认使用3306作为端口号。然后将获取到的端口号转换为整型，
		如果转换失败或者端口号不在0到65535之间，则抛出ValueError异常。

		参数:
		- port: 指定的端口号，可以为None、整数或其他可转换为整数的类型。

		返回值:
		- 验证通过后返回的整型端口号。

		异常:
		- ValueError: 当端口号无法转换为整数或不在有效范围内时抛出。
		"""
		if port is None:
			# 尝试从环境变量获取端口号，未设置则默认为3306
			port_str = os.getenv('DB_PORT', '3306')
			try:
				port = int(port_str)
			except ValueError:
				# 端口号字符串转换失败，抛出异常
				raise ValueError(f"Invalid port number: {port_str}")
			if port < 0 or port > 65535:
				# 端口号不在有效范围内，抛出异常
				raise ValueError(f"Port number must be between 0 and 65535, got: {port}")
		else:
			# 直接传入的port参数也需要进行范围验证
			if port < 0 or port > 65535:
				raise ValueError(f"Port number must be between 0 and 65535, got: {port}")
		return port

	def _init_db_connection(self):
		"""
		初始化数据库连接。

		该方法尝试生成数据库URI，创建数据库引擎，以及配置会话制作器。

		如果在这一过程中遇到InvalidRequestError错误（例如，由于无效的数据库URI），则会抛出一个
		ValueError异常，指出数据库连接初始化失败。

		返回值:
		无
		"""
		try:
			self.uri = self._generate_db_uri()  # 生成数据库URI
			self.engine = create_engine(self.uri)  # 根据URI创建数据库引擎
			self.metadata = MetaData()  # 初始化元数据
			self.Session = sessionmaker(bind=self.engine)  # 配置会话制作器
		except InvalidRequestError as e:
			# 处理无效的数据库URI
			raise ValueError("Failed to initialize database connection.") from e

	def _generate_db_uri(self) -> str:
		"""
		生成数据库的统一资源标识符（URI）。

		根据数据库类型（sqlite、mysql、mssql），构建相应的数据库连接URI。

		返回值:
			str: 数据库的URI字符串。

		异常:
			ValueError: 当数据库类型不受支持时抛出。
		"""
		if self.db_type == 'sqlite':
			# 对于sqlite数据库，构建URI格式为sqlite:///数据库文件路径
			return f'sqlite:///{self.database}'
		elif self.db_type in ['mysql', 'mssql']:
			# 对于mysql和mssql数据库，使用pymysql驱动构建URI格式为数据库类型+驱动名://用户名:密码@主机:端口/数据库名
			return f'{self.db_type}+pymysql://{self.username}:{self.password}@{self.host}:{self.port}/{self.database}'
		else:
			# 当数据库类型不在支持列表内时，抛出异常
			raise ValueError(f"Unsupported database type: {self.db_type}")

	@contextlib.contextmanager
	def managed_session(self):
		"""
		一个上下文管理器，用于自动管理数据库会话的生命周期。

		会自动开启一个数据库会话，在退出上下文管理器时，如果期间没有发生异常，则会提交这个会话的所有更改；
	 如果期间发生了异常，则会回滚所有更改，确保数据库的一致性。

		参数:
		self - 对象自身的引用。

		返回值:
		一个上下文管理器，其`yield`语句返回的值为数据库会话对象。
		"""
		session = self.Session()  # 创建数据库会话
		try:
			yield session  # 将会话对象提供给with语句块使用
			session.commit()  # 如果with块正常退出，提交会话中的所有更改
		except SQLAlchemyError as e:
			logging.error(f"Database error: {e}")  # 记录数据库错误日志
			session.rollback()  # 发生异常时，回滚会话中的所有更改
			raise  # 重新抛出捕获到的异常
		finally:
			session.close()  # 确保会话在使用后被正确关闭

	def load_table(self, table_name: str) -> pd.DataFrame:
		"""
		从数据库中加载指定的表到一个Pandas DataFrame中。

		参数:
		table_name (str): 要加载的表的名称。

		返回:
		pd.DataFrame: 从数据库表中加载的数据。
		"""
		with self.managed_session() as session:  # 管理数据库会话
			# 加载表的元数据
			table = Table(table_name, self.metadata, autoload_with=session.bind)
			query = table.select()  # 构建选择整个表的SQL查询
			# 执行查询并将结果读取到DataFrame中
			data = pd.read_sql(query, session.bind)
		return data

	def insert_data(self, table_name: str, data: Dict[str, Any]):
		"""
		将数据插入指定的数据库表中。

		参数:
		table_name (str): 要插入数据的表名。
		data (Dict[str, Any]): 要插入的数据，以字典形式提供，其中键是列名，值是要插入的值。

		返回值:
		无
		"""
		with self.managed_session() as session:  # 开始一个数据库会话
			# 根据表名加载表的元数据
			table = Table(table_name, self.metadata, autoload_with=session.bind)
			insert_statement = table.insert().values(data)  # 构造插入数据的SQL语句
			session.execute(insert_statement)  # 执行SQL语句，插入数据

	def update_data(self, table_name: str, filter_conditions: Dict[str, Any], update_values: Dict[str, Any]):
		"""
		更新指定表中的数据。

		参数:
		- table_name (str): 要更新数据的表名。
		- filter_conditions (Dict[str, Any]): 筛选条件字典，键是条件字段名，值是条件值。
		- update_values (Dict[str, Any]): 要更新的字段及其新值的字典，键是字段名，值是新值。

		返回值:
		无
		"""
		with self.managed_session() as session:  # 开始一个数据库会话
			# 根据表名加载表的元数据
			table = Table(table_name, self.metadata, autoload_with=session.bind)
			# 构造更新语句，包括更新的条件和值
			update_statement = table.update().where(text(self._dict_to_clause(filter_conditions))).values(
				update_values)
			# 执行更新语句
			session.execute(update_statement)

	def delete_data(self, table_name: str, filter_conditions: Dict[str, Any]):
		"""
		根据给定的表名和过滤条件从数据库中删除数据。

		:param table_name: 要从中删除数据的表的名称。
		:param filter_conditions: 一个字典，包含用于过滤数据的条件。键是列名，值是要匹配的值。
		"""
		with self.managed_session() as session:  # 开始一个数据库会话
			# 加载表的元数据
			table = Table(table_name, self.metadata, autoload_with=session.bind)
			# 构造一个删除语句，基于传入的过滤条件
			delete_statement = table.delete().where(text(self._dict_to_clause(filter_conditions)))
			# 执行删除语句
			session.execute(delete_statement)

	@staticmethod
	def _dict_to_clause(filter_dict: Dict[str, Any]) -> str:
		"""
		将字典转换为SQL查询条件字符串。

		参数:
		filter_dict: Dict[str, Any] - 包含要转换为查询条件的键值对的字典。
					   键代表SQL条件中的字段名，值代表条件值。

		返回值:
		str - 格式化后的SQL查询条件字符串，各个条件间通过' AND '连接。
		"""
		# 生成列表，其中每个元素是一个格式化的SQL条件字符串
		clauses = [f"{key} = '{value}'" for key, value in filter_dict.items()]
		# 将条件字符串列表用" AND "连接成一个单一的字符串
		return " AND ".join(clauses)


if __name__ == "__main__":
	db_manager = DatabaseManager(db_type=os.getenv('DB_TYPE', 'mysql'))

	# 示例：加载数据
	print("Loading data from 'users' table:")
	users_df = db_manager.load_table('users')
	print(users_df.head())

	# 示例：插入新数据
	new_user = {'id': 101, 'username': 'newuser', 'email': 'newuser@example.com'}
	db_manager.insert_data('users', new_user)
	print("New user inserted.")

	# 示例：更新数据
	update_conditions = {'id': 101}
	update_values = {'email': 'newuser_updated@example.com'}
	db_manager.update_data('users', update_conditions, update_values)
	print("User email updated.")

	# 示例：删除数据
	delete_conditions = {'id': 101}
	db_manager.delete_data('users', delete_conditions)
	print("User record deleted.")
